#include "TracerProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "VoxelProcess.h"
#include "LocalScene.h"
#include "BeamProcess.h"

void TracerProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);
	mVoxelProcess = mLocalScene->GetProcess<VoxelProcess>("VoxelProcess");
	mBeamProcess = mLocalScene->GetProcess<BeamProcess>("BeamProcess");

	mBeamEnable = true;
	mViewType = 0;


	ENQUEUE_RENDER_COMMAND(BuildModel)([this](VkCommandBuffer commandBuffer)
		{
			
			RefCountPtr<RHIBuffer> octreeBuffer = GetBuffer("OctreeProcess", "octree");
			RefCountPtr<RHIBuffer> cameraUboBuffer = GetBuffer("BeamProcess", "cameraUbo");
			RefCountPtr<RHITexture> beamImage = GetTexture("BeamProcess", "beamImage");
			mBeamSampler = RHISampler::Builder()
				.Build();


			mDescriptor = RHIDescriptor::Builder()
				.PushShader<DeferredVertexShader>()
				.PushShader<TracerFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			mDescriptor->SetValue("octree", octreeBuffer);
			mDescriptor->SetValue("cameraUbo", cameraUboBuffer);
			mDescriptor->SetValue("beamImage", beamImage, mBeamSampler);
			mDescriptorMask = mDescriptor->AllocateDescriptor();



			mColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "ColorRenderTarget");


			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void TracerProcess::OnTick(float deltaTime)
{

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this](VkCommandBuffer commandBuffer)
		{
			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, nullptr);


			PushConstants pushConstants;
			pushConstants.width = GRenderState->GetWidth();
			pushConstants.height = GRenderState->GetHeight();
			pushConstants.beamSize = mBeamProcess->mBeamSize;
			pushConstants.beamEnable = mBeamEnable ? 1 : 0;
			pushConstants.viewType = static_cast<uint32>(mViewType);


			vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(PushConstants), &pushConstants);

			vkCmdDraw(commandBuffer, 3, 1, 0, 0);

			vkCmdEndRenderPass(commandBuffer);
		});
}

void TracerProcess::OnResize()
{
	RefCountPtr<RHITexture> beamImage = GetTexture("BeamProcess", "beamImage");
	mDescriptor->SetValue("beamImage", beamImage, mBeamSampler);
	mDescriptor->UpdateDescriptor(mDescriptorMask);


	mColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "ColorRenderTarget");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}